package cn.echoparrot.domain.service;

import cn.echoparrot.domain.entity.Notebook;
import cn.echoparrot.repository.neo4j.NoteMapper;
import cn.echoparrot.repository.neo4j.NotebookRepository;

import jakarta.validation.constraints.NotNull;
import org.neo4j.driver.Driver;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.neo4j.Neo4jVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.Assert;


import java.util.Optional;

@Service
public class DefaultNotebookService implements NotebookService{

    private NotebookRepository notebookRepository;
    private NoteMapper noteMapper;
    private Driver driver;

    @Autowired
    public DefaultNotebookService(NotebookRepository notebookRepository, NoteMapper noteMapper, Driver driver) {
        this.notebookRepository = notebookRepository;
        this.noteMapper = noteMapper;
        this.driver = driver;
    }

    @Override
    @Transactional(transactionManager = "neo4jTransactionManager")
    public Optional<Notebook> find(String deviceId, Integer roleId) {
        return notebookRepository.findByDeviceIdAndRoleId(deviceId, roleId);
    }

    @Override
    @Transactional(transactionManager = "neo4jTransactionManager")
    public Notebook create(Notebook notebook) {
        Notebook.eligible(notebook);
        Notebook saved = notebookRepository.save(notebook);
        return saved;
    }


    /**
     * 打开笔记本（备忘录），可读、写、检索等
     * 必须是已经存在的
     * @param notebook
     * @param embeddingModel
     * @return
     */
    public NoteService open(Notebook notebook, EmbeddingModel embeddingModel) {
            var vectorStore = buildVectorStore(notebook, embeddingModel);
            // afterPropertiesSet 必不可少
            vectorStore.afterPropertiesSet();
            return new DefaultNoteService(notebook, embeddingModel, noteMapper, vectorStore);
    }

    private @NotNull Neo4jVectorStore buildVectorStore(Notebook notebook, EmbeddingModel embeddingModel) {
        // TODO 使用Neo4jVectorStore未必是最佳选择，但先有再优，性能需要实测。
        // neo4j Label命名不允许使用特殊字符，如'/'、'-'、'.'，但可以使用下划线'_'。
        var vectorStore = Neo4jVectorStore.builder(driver, embeddingModel)
                .databaseName("neo4j")                // Optional: defaults to "neo4j"
                .distanceType(Neo4jVectorStore.Neo4jDistanceType.COSINE) // Optional: defaults to COSINE
                .embeddingDimension(embeddingModel.dimensions())                      // Optional: defaults to 1536
                .label(notebook.labelName())                     // Optional: defaults to "Document"
                .constraintName(notebook.constraintName())
                .embeddingProperty("embedding")        // Optional: defaults to "embedding"
                .indexName(notebook.indexName())             // Optional: defaults to "spring-ai-document-index"
                .initializeSchema(true)                // Optional: defaults to false
                .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy
                //.textProperty("content")
                .build();
        return vectorStore;
    }
}
